import logging
import numpy as np
import pandas as pd
from qtorch.core.qstrategy import QStrategy

class EnhancedRSIStrategy(QStrategy):
    """
    增强版RSI策略，提供多种配置选项和过滤条件
    
    参数:
        period: RSI计算周期
        overbought: 超买阈值
        oversold: 超卖阈值
        use_ma_filter: 是否使用均线过滤
        ma_period: 均线周期
        ma_type: 均线类型 ('sma'或'ema')
        use_trend_filter: 是否使用趋势过滤
        trend_period: 趋势判断周期
        exit_rsi_overbought: 平多仓RSI阈值
        exit_rsi_oversold: 平空仓RSI阈值
        use_fixed_exit: 是否使用固定退出（而非反向信号退出）
    """
    
    def __init__(self, period=14, overbought=70, oversold=30, 
                use_ma_filter=False, ma_period=50, ma_type='sma',
                use_trend_filter=False, trend_period=100,
                exit_rsi_overbought=65, exit_rsi_oversold=35,
                use_fixed_exit=False):
        """初始化RSI策略"""
        super().__init__()
        self.period = period
        self.overbought = overbought
        self.oversold = oversold
        self.use_ma_filter = use_ma_filter
        self.ma_period = ma_period
        self.ma_type = ma_type
        self.use_trend_filter = use_trend_filter
        self.trend_period = trend_period
        self.exit_rsi_overbought = exit_rsi_overbought
        self.exit_rsi_oversold = exit_rsi_oversold
        self.use_fixed_exit = use_fixed_exit
        
    def initialize(self):
        """策略初始化"""
        logging.info(f"初始化增强RSI策略: 周期={self.period}, 超买={self.overbought}, 超卖={self.oversold}")
        if self.use_ma_filter:
            logging.info(f"使用均线过滤: 类型={self.ma_type}, 周期={self.ma_period}")
        if self.use_trend_filter:
            logging.info(f"使用趋势过滤: 周期={self.trend_period}")
            
    def next(self):
        """
        处理当前数据点，生成交易信号
        
        返回:
            int: 1(买入)、0(持有)、-1(卖出)
        """
        # 计算RSI
        rsi = self.RSI(self.period)
        if rsi is None:
            return 0  # 数据不足
        
        # 获取当前价格
        if 'close' not in self.data:
            return 0  # 数据缺失
            
        price = self.data['close']
        
        # 计算均线（如果启用均线过滤）
        if self.use_ma_filter:
            if self.ma_type.lower() == 'ema':
                ma = self.EMA(self.ma_period)
            else:
                ma = self.SMA(self.ma_period)
                
            if ma is None:
                return 0  # 数据不足
        
        # 计算趋势（如果启用趋势过滤）
        if self.use_trend_filter:
            if len(self.prices) < self.trend_period:
                return 0  # 数据不足
                
            trend_ma = self.SMA(self.trend_period)
            if trend_ma is None:
                return 0  # 数据不足
                
            trend = 1 if price > trend_ma else -1
        
        # 生成交易信号
        signal = 0
        
        # 超卖区域，考虑买入
        if rsi < self.oversold and self.position <= 0:
            # 应用过滤条件
            can_buy = True
            
            if self.use_ma_filter:
                can_buy = can_buy and price > ma
                
            if self.use_trend_filter:
                can_buy = can_buy and trend > 0
                
            if can_buy:
                signal = 1  # 买入信号
        
        # 超买区域，考虑卖出
        elif rsi > self.overbought and self.position >= 0:
            # 应用过滤条件
            can_sell = True
            
            if self.use_ma_filter:
                can_sell = can_sell and price < ma
                
            if self.use_trend_filter:
                can_sell = can_sell and trend < 0
                
            if can_sell:
                signal = -1  # 卖出信号
                
        # 固定退出（如果启用）
        elif self.use_fixed_exit:
            if self.position > 0 and rsi > self.exit_rsi_overbought:
                signal = -1  # 平多仓
            elif self.position < 0 and rsi < self.exit_rsi_oversold:
                signal = 1   # 平空仓
        
        return signal